import numpy as np
import cv2


def boxcount_q(Z, k, q=0):
    if q == 0:
        S = np.add.reduceat(
            np.add.reduceat(Z, np.arange(0, Z.shape[0], k), axis=0),
            np.arange(0, Z.shape[1], k), axis=1)
        return np.count_nonzero(S)
    else:
        I = np.ones_like(Z)
        I = np.add.reduceat(
            np.add.reduceat(I, np.arange(0, I.shape[0], k), axis=0),
            np.arange(0, I.shape[1], k), axis=1)
        S = np.add.reduceat(
            np.add.reduceat(Z, np.arange(0, Z.shape[0], k), axis=0),
            np.arange(0, Z.shape[1], k), axis=1)
        S = (S.astype(np.float32) / I) ** q
    return S.sum()


def multifractal_dimension(Z, q=0):
    '''
    输入：
        Z是BW图像，只有0和1两种数值
        q可以是单个数值，也可以是list
    输出：
        如果q是单个数值，那么输出就是单个数值
        如果q是list，那么输出就是numpy的数组
    '''

    assert(len(Z.shape) == 2)

    p = min(Z.shape)
    n = 2 ** np.floor(np.log(p) / np.log(2))
    n = int(np.log(n) / np.log(2))
    sizes = 2 ** np.arange(n, 0, -1)

    if isinstance(q, list):
        assert 1 not in q
        # 其实这里可以加速。为了简洁，先这样写。
        result = []
        for single_q in q:
            counts = []
            for size in sizes:
                counts.append(boxcount_q(Z, size, single_q))
            coeffs = np.polyfit(np.log(sizes), np.log(counts), 1) / (1 - single_q)
            result.append(-coeffs[0])
        return np.array(result)
    else:
        assert(q != 1)
        counts = []
        for size in sizes:
            counts.append(boxcount_q(Z, size, q))
        # Fit
        coeffs = np.polyfit(np.log(sizes), np.log(counts), 1) / (1 - q)
        return -coeffs[0]
